!mkdir StanfordDogs !wget http://vision.stanford.edu/aditya86/ImageNetDogs/images.tar && tar -xf images.tar -C StanfordDogs !wget http://vision.stanford.edu/aditya86/ImageNetDogs/annotation.tar && tar -xf annotation.tar -C StanfordDogs
# Create the directory tree !mkdir StanfordDogs_final StanfordDogs_final/Images StanfordDogs_final/Annotation !mkdir StanfordDogs_final/Images/Retriever StanfordDogs_final/Annotation/Retriever !mkdir StanfordDogs_final/Images/Levrier StanfordDogs_final/Annotation/Levrier !mkdir StanfordDogs_final/Images/Terrier StanfordDogs_final/Annotation/Terrier !mkdir StanfordDogs_final/Images/Bull StanfordDogs_final/Annotation/Bull
# Retriever !cp -r StanfordDogs/Images/n02099267-flat-coated_retriever/* StanfordDogs_final/Images/Retriever !cp -r StanfordDogs/Annotation/n02099267-flat-coated_retriever/* StanfordDogs_final/Annotation/Retriever !cp -r StanfordDogs/Images/n02099601-golden_retriever/* StanfordDogs_final/Images/Retriever !cp -r StanfordDogs/Annotation/n02099601-golden_retriever/* StanfordDogs_final/Annotation/Retriever !cp -r StanfordDogs/Images/n02099712-Labrador_retriever/* StanfordDogs_final/Images/Retriever !cp -r StanfordDogs/Annotation/n02099712-Labrador_retriever/* StanfordDogs_final/Annotation/Retriever !cp -r StanfordDogs/Images/n02099849-Chesapeake_Bay_retriever/* StanfordDogs_final/Images/Retriever !cp -r StanfordDogs/Annotation/n02099849-Chesapeake_Bay_retriever/* StanfordDogs_final/Annotation/Retriever # Levrier !cp -r StanfordDogs/Images/n02090622-borzoi/* StanfordDogs_final/Images/Levrier !cp -r StanfordDogs/Annotation/n02090622-borzoi/* StanfordDogs_final/Annotation/Levrier !cp -r StanfordDogs/Images/n02091032-Italian_greyhound/* StanfordDogs_final/Images/Levrier !cp -r StanfordDogs/Annotation/n02091032-Italian_greyhound/* StanfordDogs_final/Annotation/Levrier !cp -r StanfordDogs/Images/n02091134-whippet/* StanfordDogs_final/Images/Levrier !cp -r StanfordDogs/Annotation/n02091134-whippet/* StanfordDogs_final/Annotation/Levrier !cp -r StanfordDogs/Images/n02091831-Saluki/* StanfordDogs_final/Images/Levrier !cp -r StanfordDogs/Annotation/n02091831-Saluki/* StanfordDogs_final/Annotation/Levrier # Terrier !cp -r StanfordDogs/Images/n02096177-cairn/* StanfordDogs_final/Images/Terrier !cp -r StanfordDogs/Annotation/n02096177-cairn/* StanfordDogs_final/Annotation/Terrier !cp -r StanfordDogs/Images/n02096294-Australian_terrier/* StanfordDogs_final/Images/Terrier !cp -r StanfordDogs/Annotation/n02096294-Australian_terrier/* StanfordDogs_final/Annotation/Terrier !cp -r StanfordDogs/Images/n02094114-Norfolk_terrier/* StanfordDogs_final/Images/Terrier !cp -r StanfordDogs/Annotation/n02094114-Norfolk_terrier/* StanfordDogs_final/Annotation/Terrier !cp -r StanfordDogs/Images/n02094258-Norwich_terrier/* StanfordDogs_final/Images/Terrier !cp -r StanfordDogs/Annotation/n02094258-Norwich_terrier/* StanfordDogs_final/Annotation/Terrier # Bull !cp -r StanfordDogs/Images/n02096585-Boston_bull/* StanfordDogs_final/Images/Bull !cp -r StanfordDogs/Annotation/n02096585-Boston_bull/* StanfordDogs_final/Annotation/Bull !cp -r StanfordDogs/Images/n02108915-French_bulldog/* StanfordDogs_final/Images/Bull !cp -r StanfordDogs/Annotation/n02108915-French_bulldog/* StanfordDogs_final/Annotation/Bull !cp -r StanfordDogs/Images/n02110958-pug/* StanfordDogs_final/Images/Bull !cp -r StanfordDogs/Annotation/n02110958-pug/* StanfordDogs_final/Annotation/Bull !cp -r StanfordDogs/Images/n02112706-Brabancon_griffon/* StanfordDogs_final/Images/Bull !cp -r StanfordDogs/Annotation/n02112706-Brabancon_griffon/* StanfordDogs_final/Annotation/Bull
# Number of examples by class ! ls -1 StanfordDogs_final/Images/Retriever | wc -l ! ls -1 StanfordDogs_final/Images/Levrier | wc -l ! ls -1 StanfordDogs_final/Images/Terrier | wc -l ! ls -1 StanfordDogs_final/Images/Bull | wc -l
import os
import io
import cv2
import time
import copy
import torch
import contextlib
import torchvision
import torch.utils.data
import numpy as np
import torchvision.transforms as T
from PIL import Image
# from tqdm import tqdm
from tqdm.notebook import tqdm
from pycocotools.coco import COCO
from matplotlib import pyplot as plt
from pycocotools.cocoeval import COCOeval
from pycocotools import mask as coco_mask
from torch.utils.data.dataset import Dataset
from torch.utils.data.dataloader import DataLoader
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights
from torch.utils.data import Subset
import xml.etree.ElementTree as ET
from torchvision.transforms.functional import to_pil_image
from torch.utils.data import random_split
# Seed.
np.random.seed(66)
torch.manual_seed(66)
DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print(DEVICE)model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.to(DEVICE)
transform = T.Compose([T.ToTensor()])
def show_image_with_boxes(image, boxes, labels, scores, title, threshold=0.7):
image = np.array(image)
for box, label, score in zip(boxes, labels, scores):
if score < threshold:
continue
xmin, ymin, xmax, ymax = box.astype(int)
cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
cv2.putText(image, f"{label}:{score:.2f}", (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
plt.figure(figsize=(12,8))
plt.imshow(image)
plt.title(title)
plt.axis('off')
plt.show()
img_path = "/content/StanfordDogs_final/Images/Bull/n02096585_1143.jpg"
img = Image.open(img_path).convert("RGB")
transform = T.Compose([T.ToTensor()])
img_t = transform(img).to(DEVICE)
# Prediction
model.eval()
with torch.no_grad():
pred = model([img_t])[0]
boxes = pred['boxes'].cpu().numpy()
labels = pred['labels'].cpu().numpy()
scores = pred['scores'].cpu().numpy()
COCO_INSTANCE_CATEGORY_NAMES = [
'__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'street sign',
'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'hat', 'backpack',
'umbrella', 'shoe', 'eye glasses', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
'skateboard', 'surfboard', 'tennis racket', 'bottle', 'plate', 'wine glass',
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
'couch', 'potted plant', 'bed', 'mirror', 'dining table', 'window', 'desk',
'toilet', 'door', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'blender', 'book',
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush',
'hair brush'
]
# Print with the name of each present class
labels_names = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in labels]
show_image_with_boxes(img, boxes, labels_names, scores, "Test Faster RCNN")print(img.size)
# Transformation : resize + ToTensor
transform = T.Compose([
T.Resize((250, 250)),
T.ToTensor()
])
img_t = transform(img) # Tensor C,H,W
print(img_t.shape)
img_np = img_t.permute(1, 2, 0).numpy()
plt.figure(figsize=(6, 12))
plt.imshow(img_np)
plt.axis("off")
plt.show()class StanfordDogsDataset(Dataset):
def __init__(self, root, transforms):
"""
root: dossier racine du dataset Stanford Dogs
root/Images -> images .jpg
root/Annotation -> annotations XML
"""
self.root = root
self.transforms = transforms
self.img = []
for subdir, _, files in os.walk(os.path.join(root, "Images")):
for f in files:
if f.lower().endswith(".jpg"):
self.img.append(os.path.join(subdir, f))
self.img = sorted(self.img)
self.ann = []
for subdir, _, files in os.walk(os.path.join(root, "Annotation")):
for f in files:
if f.lower():
self.ann.append(os.path.join(subdir, f))
self.ann = sorted(self.ann)
image_root = os.path.join(root, "Images")
# Mapping
self.classes = sorted(os.listdir(image_root))
self.class_to_idx = {cls_name: idx+1 for idx, cls_name in enumerate(self.classes)}
# 0 is for background
self.group_mapping = {
# Bull
"Boston_bull": 1, "French_bulldog": 1, "pug": 1, "Brabancon_griffon": 1,
# Levrier
"borzoi": 2, "Italian_greyhound": 2, "whippet": 2, "Saluki": 2,
# Retriever
"flat-coated_retriever": 3, "golden_retriever": 3, "Labrador_retriever": 3, "Chesapeake_Bay_retriever": 3,
# Terrier
"cairn": 4, "Australian_terrier": 4, "Norfolk_terrier": 4, "Norwich_terrier": 4,
}
def __len__(self):
return len(self.img)
def __getitem__(self, idx):
img_path = os.path.join(self.root, "Images", self.img[idx])
ann_path = os.path.join(self.root, "Annotation", self.ann[idx])
img = Image.open(img_path).convert("RGB")
# Parse XML
tree = ET.parse(ann_path)
root_xml = tree.getroot()
boxes = []
labels = []
iscrowd = []
for obj in root_xml.findall("object"):
bndbox = obj.find("bndbox")
xmin = int(bndbox.find("xmin").text)
ymin = int(bndbox.find("ymin").text)
xmax = int(bndbox.find("xmax").text)
ymax = int(bndbox.find("ymax").text)
boxes.append([xmin, ymin, xmax, ymax])
class_name = obj.find("name").text.strip()
labels.append(self.group_mapping[class_name])
iscrowd.append(0)
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(labels, dtype=torch.int64)
iscrowd = torch.as_tensor(iscrowd, dtype=torch.int64)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
target = {
"boxes": boxes,
"labels": labels,
"image_id": image_id,
"area": area,
"iscrowd": iscrowd
}
img = self.transforms(img)
return img, targettransforms = T.Compose([T.ToTensor()])
dataset = StanfordDogsDataset("/content/StanfordDogs_final", transforms)print(dataset.class_to_idx)
def collate_fn(batch):
images = [item[0] for item in batch]
targets = [item[1] for item in batch]
return images, targets
n_total = len(dataset.img)
val_ratio = 0.1
test_ratio = 0.1
n_val = int(n_total * val_ratio)
n_test = int(n_total * test_ratio)
n_train = n_total - n_val - n_test
# Random split
train_dataset, val_dataset, test_dataset = random_split(dataset, [n_train, n_val, n_test])
print(f"Total: {len(dataset)}, Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")# Hyperparametres
EPOCHS = 10
BATCH_SIZE = 64
# DataLoaders
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
# Model
weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT # pre-trained on COCO
model = fasterrcnn_resnet50_fpn(weights=weights)
num_classes = len(dataset.classes) + 1
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # New head of the model
model.transform.min_size = (240,) # Force to downsize images shape
model.transform.max_size = 260
model.to(DEVICE)
for param in model.parameters():
param.requires_grad = False
# Unfreeze layer4
for name, param in model.backbone.body.named_parameters():
if "layer4" in name:
param.requires_grad = True
# Unfreeze roi_heads
for param in model.roi_heads.parameters():
param.requires_grad = True
params_to_optimize = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params_to_optimize, lr=0.005, momentum=0.9, weight_decay=0.0005)
train_losses = []
val_losses = []
scaler = torch.amp.GradScaler()
for epoch in range(EPOCHS):
model.train()
running_loss = 0.0
tepoch = tqdm(enumerate(train_loader), total=len(train_loader), unit="batch")
tepoch.set_description(f"Epoch {epoch+1}/{EPOCHS} [Train]")
for i, (images, targets) in tepoch:
images = [img.to(DEVICE) for img in images]
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]
with torch.amp.autocast(device_type='cuda'):
loss_dict = model(images, targets)
loss = sum(loss for loss in loss_dict.values())
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
running_loss += loss.item()
tepoch.set_postfix(loss=loss.item())
epoch_train_loss = running_loss / len(train_loader)
train_losses.append(epoch_train_loss)
# Validation
# No eval mode for this model to get validation loss
# So Dropout and BatchNorm remain... but the loss curve is still useful
running_val_loss = 0.0
tepoch_val = tqdm(enumerate(val_loader), total=len(val_loader), unit="batch")
tepoch_val.set_description(f"Epoch {epoch+1}/{EPOCHS} [Val]")
with torch.no_grad():
for i, (images, targets) in tepoch_val:
images = [img.to(DEVICE) for img in images]
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]
with torch.amp.autocast(device_type='cuda'):
loss_dict = model(images, targets)
loss = sum(loss for loss in loss_dict.values())
running_val_loss += loss.item()
tepoch_val.set_postfix(loss=loss.item())
epoch_val_loss = running_val_loss / len(val_loader)
val_losses.append(epoch_val_loss)
print(f"Epoch {epoch+1}/{EPOCHS} - Train Loss: {epoch_train_loss:.4f} - Val Loss: {epoch_val_loss:.4f}")0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
0%| | 0/36 [00:00<?, ?batch/s]
0%| | 0/5 [00:00<?, ?batch/s]
# Loss curves
plt.figure(figsize=(8,5))
plt.plot(range(1, len(train_losses)+1), train_losses, label='Train Loss')
plt.plot(range(1, len(val_losses)+1), val_losses, label='Validation Loss')
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training & Validation Loss")
plt.legend()
plt.show()class CocoEvaluator:
def __init__(self, coco_gt, iou_types):
assert isinstance(iou_types, (list, tuple))
coco_gt = copy.deepcopy(coco_gt)
self.coco_gt = coco_gt
self.iou_types = iou_types
self.coco_eval = {}
for iou_type in iou_types:
self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)
self.img_ids = []
self.eval_imgs = {k: [] for k in iou_types}
def update(self, predictions):
img_ids = list(np.unique(list(predictions.keys())))
self.img_ids.extend(img_ids)
for iou_type in self.iou_types:
results = self.prepare(predictions, iou_type)
with contextlib.redirect_stdout(io.StringIO()):
coco_dt = COCO.loadRes(self.coco_gt, results) if results else COCO()
coco_eval = self.coco_eval[iou_type]
coco_eval.cocoDt = coco_dt
coco_eval.params.imgIds = list(img_ids)
img_ids, eval_imgs = coco_evaluate(coco_eval)
self.eval_imgs[iou_type].append(eval_imgs)
def synchronize_between_processes(self):
for iou_type in self.iou_types:
self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])
def accumulate(self):
for coco_eval in self.coco_eval.values():
coco_eval.accumulate()
def summarize(self):
for iou_type, coco_eval in self.coco_eval.items():
print(f"IoU metric: {iou_type}")
coco_eval.summarize()
def prepare(self, predictions, iou_type):
if iou_type == "bbox":
return self.prepare_for_coco_detection(predictions)
if iou_type == "segm":
return self.prepare_for_coco_segmentation(predictions)
if iou_type == "keypoints":
return self.prepare_for_coco_keypoint(predictions)
raise ValueError(f"Unknown iou type {iou_type}")
def prepare_for_coco_detection(self, predictions):
coco_results = []
for original_id, prediction in predictions.items():
if len(prediction) == 0:
continue
boxes = prediction["boxes"]
boxes = convert_to_xywh(boxes).tolist()
scores = prediction["scores"].tolist()
labels = prediction["labels"].tolist()
coco_results.extend(
[
{
"image_id": original_id,
"category_id": labels[k],
"bbox": box,
"score": scores[k],
}
for k, box in enumerate(boxes)
]
)
return coco_results
def convert_to_xywh(boxes):
xmin, ymin, xmax, ymax = boxes.unbind(1)
return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
def merge(img_ids, eval_imgs):
import torch.distributed as dist
def is_dist_avail_and_initialized():
if not dist.is_available():
return False
if not dist.is_initialized():
return False
return True
def get_world_size():
if not is_dist_avail_and_initialized():
return 1
return dist.get_world_size()
def all_gather(data):
world_size = get_world_size()
if world_size == 1:
return [data]
data_list = [None] * world_size
dist.all_gather_object(data_list, data)
return data_list
all_img_ids = all_gather(img_ids)
all_eval_imgs = all_gather(eval_imgs)
merged_img_ids = []
for p in all_img_ids:
merged_img_ids.extend(p)
merged_eval_imgs = []
for p in all_eval_imgs:
merged_eval_imgs.append(p)
merged_img_ids = np.array(merged_img_ids)
merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
# Keep only unique (and in sorted order) images.
merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
merged_eval_imgs = merged_eval_imgs[..., idx]
return merged_img_ids, merged_eval_imgs
def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
img_ids, eval_imgs = merge(img_ids, eval_imgs)
img_ids = list(img_ids)
eval_imgs = list(eval_imgs.flatten())
coco_eval.evalImgs = eval_imgs
coco_eval.params.imgIds = img_ids
coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
def coco_evaluate(imgs):
with contextlib.redirect_stdout(io.StringIO()):
imgs.evaluate()
return imgs.params.imgIds, np.asarray(imgs.evalImgs).reshape(-1, len(imgs.params.areaRng), len(imgs.params.imgIds))def convert_to_coco_api(ds):
coco_ds = COCO()
# Annotation IDs need to start at 1, not 0, see torchvision issue #1530.
ann_id = 1
dataset = {"images": [], "categories": [], "annotations": [],
# Add a placeholder 'info' dictionary
"info": {"description": "Converted dataset", "version": "1.0", "year": 2023,
"date_created": "2023/01/01"}}
categories = set()
for img_idx in range(len(ds)):
# Find better way to get target.
# targets = ds.get_annotations(img_idx)
img, targets = ds[img_idx]
image_id = targets["image_id"].item()
img_dict = {}
img_dict["id"] = image_id
img_dict["height"] = img.shape[-2]
img_dict["width"] = img.shape[-1]
dataset["images"].append(img_dict)
bboxes = targets["boxes"]
bboxes[:, 2:] -= bboxes[:, :2]
bboxes = bboxes.tolist()
labels = targets["labels"].tolist()
areas = targets["area"].tolist()
iscrowd = targets["iscrowd"].tolist()
num_objs = len(bboxes)
for i in range(num_objs):
ann = {}
ann["image_id"] = image_id
ann["bbox"] = bboxes[i]
ann["category_id"] = labels[i]
categories.add(labels[i])
ann["area"] = areas[i]
ann["iscrowd"] = iscrowd[i]
ann["id"] = ann_id
dataset["annotations"].append(ann)
ann_id += 1
dataset["categories"] = [{"id": i} for i in sorted(categories)]
coco_ds.dataset = dataset
coco_ds.createIndex()
return coco_ds
def get_coco_api_from_dataset(dataset):
for _ in range(10):
if isinstance(dataset, torchvision.datasets.CocoDetection):
break
if isinstance(dataset, torch.utils.data.Subset):
dataset = dataset.dataset
if isinstance(dataset, torchvision.datasets.CocoDetection):
return dataset.coco
return convert_to_coco_api(dataset)def evaluate(model, data_loader, device):
n_threads = torch.get_num_threads()
torch.set_num_threads(1)
cpu_device = torch.device("cpu")
coco = get_coco_api_from_dataset(data_loader.dataset)
coco_evaluator = CocoEvaluator(coco, ["bbox"])
model.eval()
for images, targets in data_loader:
images = list(img.to(device) for img in images)
if torch.cuda.is_available():
torch.cuda.synchronize()
model_time = time.time()
with torch.no_grad():
outputs = model(images)
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
model_time = time.time() - model_time
res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_time = time.time() - evaluator_time
coco_evaluator.synchronize_between_processes()
# Accumulate predictions from all images.
coco_evaluator.accumulate()
coco_evaluator.summarize()
torch.set_num_threads(n_threads)
return coco_evaluator
evaluate(model, test_loader, DEVICE)<__main__.CocoEvaluator at 0x78fdc9858860>
model.eval()
myClasses = ['__background__', 'Bull', 'Levrier', 'Retreiver', 'Terrier']
for i in range(10):
img, target = test_dataset[i]
img = img.to(DEVICE)
# Prediction
with torch.no_grad():
pred = model([img])[0]
boxes = pred['boxes'].cpu().numpy()
labels = pred['labels'].cpu().numpy()
scores = pred['scores'].cpu().numpy()
labels_names = [myClasses[i] for i in labels]
if len(scores) == 0:
continue # No prediction
# Only print the most likely box
best_idx = scores.argmax()
best_box = boxes[best_idx:best_idx+1]
best_label = [myClasses[labels[best_idx]]]
best_score = [scores[best_idx]]
img_display = to_pil_image(img)
show_image_with_boxes(img_display, best_box, best_label, best_score,
f"Vrai label : {myClasses[target['labels'][0]]}, Prediction : {best_label[0]}", threshold=0.0)
plt.show()